import sys
sys.path.append("")
import re
import numpy as np
from bert4keras.backend import  keras,K
from bert4keras.models import  build_transformer_model
from bert4keras.tokenizers import Tokenizer,load_vocab
from bert4keras.optimizers import Adam
from bert4keras.snippets import sequence_padding,open
from bert4keras.snippets import DataGenerator,AutoRegressiveDecoder
from keras.layers import Input,Embedding,Reshape

# 模型配置
maxlen = 128
batch_size = 32
num_classes = 2
epochs = 20

import os
if 'win' in sys.platform:
    BertPath = r'E:\nlp-data\bert\model_chinese'
else:
    BertPath = r'/data/nlp_data/bert/model_chinese/'

config_path = os.path.join(BertPath,'bert_config.json')
checkpoint_path = os.path.join(BertPath,'bert_model.ckpt')
dict_path = os.path.join(BertPath,'vocab.txt')

token_dict,keep_tokens = load_vocab(
    dict_path=dict_path,
    simplified=True,
    startwith=['[PAD]', '[UNK]', '[CLS]', '[SEP]'],
)
tokenizer = Tokenizer(token_dict,do_lower_case=True)

def load_data(filenames):
    D = []
    for filename in filenames:
        with open(filename, encoding='utf-8') as f:
            for l in f:
                text, label = l.strip().split('\t')
                if len(text) <= maxlen - 2:
                    D.append((text, int(label)))
                else:
                    tmp = ''
                    for t in re.findall(u'.*?[。！\n]', text):
                        if tmp and len(tmp) + len(t) > maxlen - 2:
                            D.append((tmp, int(label)))
                            tmp = ''
                        tmp += t
                    if tmp:
                        D.append((tmp, int(label)))
    return D

# 加载数据集
data = load_data([
    'datasets/sentiment/sentiment.train.data',
    'datasets/sentiment/sentiment.valid.data',
    'datasets/sentiment/sentiment.test.data',
])

class data_generator(DataGenerator):
    def __iter__(self,random=False):
        batch_token_ids,batch_segment_ids,batch_labels = [],[],[]
        for is_end,(text,label) in self.sample(random):
            token_ids,segment_ids = tokenizer.encode(text,max_length=maxlen)
            batch_token_ids.append(token_ids)
            batch_segment_ids.append(segment_ids)
            batch_labels.append([label])
            if len(batch_token_ids) == self.batch_size or is_end:
                batch_token_ids = sequence_padding(batch_token_ids)
                batch_segment_ids = sequence_padding(batch_segment_ids)
                batch_labels = sequence_padding(batch_labels)
                yield [batch_token_ids,batch_segment_ids,batch_labels],None
                batch_token_ids,batch_segment_ids,batch_labels= [],[],[]


c_in = Input(shape=(1,))
c = Embedding(2,128)(c_in)
c = Reshape((128,))(c)

model = build_transformer_model(
    config_path,
    checkpoint_path,
    application='lm',
    keep_tokens=keep_tokens,
    layer_norm_cond=c,
    additional_input_layers=c_in,
)

model.summary()

y_true = model.input[0][:,1:]
y_mask = model.get_layer('Embedding-Token').output_mask[:,1:]
y_mask = K.cast(y_mask,K.floatx())
y_pred = model.output[:,:-1]
cross_entropy = K.sparse_categorical_crossentropy(y_true,y_pred)
cross_entropy= K.sum(cross_entropy * y_mask) /K.sum(y_mask)

model.add_loss(cross_entropy)
model.compile(optimizer=Adam(1e-5))

class RandomSentiment(AutoRegressiveDecoder):
    """根据情感标签（0:负，1:正）随机生成一批句子
        """
    @AutoRegressiveDecoder.set_rtype('probas')
    def predict(self, inputs, output_ids, step):
        token_ids = output_ids
        segmetn_ids = np.zeros_like(token_ids)
        return model.predict([token_ids,segmetn_ids,inputs[0]])[:,-1]

    def generate(self,label,n=1,topk=5):
        results = self.random_sample([[label]],n,topk)
        return [tokenizer.decode(ids) for ids in results]


random_sentiment = RandomSentiment(
    start_id = tokenizer._token_start_id,
    end_id = tokenizer._token_end_id,
    maxlen=maxlen
)

def just_show():
    print(u'正面采样：')
    print(random_sentiment.generate(1,5,5),'\n')
    print(u"负面采样:")
    print(random_sentiment.generate(0,5,5),'\n')

class Evaluate(keras.callbacks.Callback):
    def __init__(self):
        self.lowest = 1e10

    def on_epoch_end(self, epoch, logs=None):
        if logs["loss"] <= self.lowest:
            self.lowest = logs['loss']
            model.save_weights("output/conditional_language.weights")

        just_show()
if __name__ == '__main__':

    if not os.path.exists("output/conditional_language.weights"):
        evaluator = Evaluate()
        train_generator = data_generator(data, batch_size)

        model.fit_generator(
            train_generator.forfit(),
            steps_per_epoch=len(train_generator),
            epochs=epochs,
            callbacks=[evaluator]
        )
    else:
        model.load_weights("output/conditional_language.weights")
        just_show()





